from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction


def calculate_exact_match(predicted, actual):
    """计算精确匹配得分"""
    return 1.0 if str(predicted).strip() == str(actual).strip() else 0.0


def calculate_bleu_score(predicted, actual):
    """计算BLEU得分"""
    # 将代码转换为token列表
    pred_tokens = str(predicted).strip().split()
    actual_tokens = str(actual).strip().split()

    # 处理空字符串情况
    if not pred_tokens or not actual_tokens:
        return 0.0

    # 计算BLEU-4得分
    smoothing = SmoothingFunction().method1
    bleu_score = sentence_bleu(
        [actual_tokens],
        pred_tokens,
        weights=(0.25, 0.25, 0.25, 0.25),
        smoothing_function=smoothing
    )
    return bleu_score


def calculate_codebleu_score(predicted, actual):
    """
    计算CodeBLEU得分（简化版实现）
    注意：这是简化版本，实际应用中建议使用专门的CodeBLEU库
    """
    try:
        # 1. n-gram match (BLEU component)
        bleu_component = calculate_bleu_score(predicted, actual)

        # 2. AST match (简化处理，基于语法结构关键词)
        pred_str = str(predicted)
        actual_str = str(actual)

        # 提取关键词作为AST节点的近似表示
        keywords = ['if', 'else', 'for', 'while', 'class', 'def', 'return', 'import', 'try', 'except']
        pred_keywords = [kw for kw in keywords if kw in pred_str]
        actual_keywords = [kw for kw in keywords if kw in actual_str]

        if len(actual_keywords) > 0:
            ast_match = len(set(pred_keywords) & set(actual_keywords)) / len(set(actual_keywords))
        else:
            ast_match = 1.0 if len(pred_keywords) == 0 else 0.0

        # 3. Data-flow match (简化处理)
        # 这里简单地基于变量名匹配
        import re
        pred_vars = set(re.findall(r'[a-zA-Z_][a-zA-Z0-9_]*', pred_str))
        actual_vars = set(re.findall(r'[a-zA-Z_][a-zA-Z0-9_]*', actual_str))

        if len(actual_vars) > 0:
            dataflow_match = len(pred_vars & actual_vars) / len(actual_vars)
        else:
            dataflow_match = 1.0 if len(pred_vars) == 0 else 0.0

        # 组合各组件得分 (权重可调整)
        codebleu_score = 0.3 * bleu_component + 0.3 * ast_match + 0.4 * dataflow_match
        return codebleu_score

    except Exception as e:
        # 出现错误时回退到普通BLEU
        return calculate_bleu_score(predicted, actual)
